# -*- coding: utf-8 -*-

"""
Created on Thu Jun 30 09:46:58 2022

@author: nikoreed
"""

import numpy as np
import mat73

file_dir='/'
save_dir='/'

num_files=80; #highest MAT file, +1 from highest npy
start_ind=0 #npy start index
offset=0 #difference between MATLAB file-1 and desired npy file number (for combining multiple batches of data)

for nf in range(start_ind,num_files):
    print(nf)
    
    batch_str='TrainingDat_64'
    save_str='64x64_50um'
    
    train_out_file=file_dir+batch_str+'_out_'+str(nf+1).zfill(3)+'.mat'
    train_in_file=file_dir+batch_str+'_in_'+str(nf+1).zfill(3)+'.mat'

    train_in_dict = mat73.loadmat(train_in_file)
    train_out_dict = mat73.loadmat(train_out_file)
    
    train_in_key = list(train_in_dict.keys())[0]
    train_out_key = list(train_out_dict.keys())[0]
    
    
    data_in = np.empty(shape=train_in_dict[train_in_key]['Bx'].shape + (4,), dtype='f') ## added data type
    data_out = np.empty(shape=train_out_dict[train_out_key]['Jx'].shape + (2,), dtype='f') ## added data type
    
    data_in[..., 0] = train_in_dict[train_in_key]['Bx']
    data_in[..., 1] = train_in_dict[train_in_key]['By']
    data_in[..., 2] = train_in_dict[train_in_key]['Bz']
    standoff_distances = train_in_dict[train_in_key]['z'] 
    
    data_out[..., 0] = train_out_dict[train_out_key]['Jx']
    data_out[..., 1] = train_out_dict[train_out_key]['Jy']
    
    data_in=np.nan_to_num(data_in)
    data_out=np.nan_to_num(data_out)
    
    data_in=data_in
    
    file_in=save_dir+save_str+'_Bxyz_'+str(nf+offset).zfill(3)+'.npy'
    file_out=save_dir+save_str+'_Jxy_'+str(nf+offset).zfill(3)+'.npy'
    file_z = save_dir+'standoff_distances_'+str(nf+offset).zfill(3)+'.npy'
    
    np.save(file_in,data_in)
    np.save(file_out,data_out)
    np.save(file_z,standoff_distances)